import time

# import os
# import django
# # 设置 DJANGO_SETTINGS_MODULE 环境变量（引入settings文件）
# os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'ROOT.settings')
# # 加载 Django 项目配置
# django.setup()
# from apps.ipo.models import Feature

import numpy as np
from tensorflow.keras.applications import ResNet50, ResNet101
from tensorflow.keras.applications.resnet import preprocess_input  # 注意：ResNet101 的预处理模块
from tensorflow.keras.preprocessing import image

import multiprocessing
from functools import partial

import os
import random
import time



# 多进程初始化函数
def init_model(model_name):
    global model
    if model_name == 'ResNet101':
        model = ResNet101(weights='imagenet', include_top=False, pooling='avg')
    elif model_name == 'ResNet50':
        model = ResNet50(weights='imagenet', include_top=False, pooling='avg')


def _extract_single_feature(img_path, model_type):
    try:
        # 加载图片并调整大小为224x224（ResNet模型的标准输入尺寸）
        img = image.load_img(img_path, target_size=(224, 224))
        # 将图片转换为numpy数组
        img_data = image.img_to_array(img)
        # 扩展数组维度，添加batch维度（模型需要batch输入）
        img_data = np.expand_dims(img_data, axis=0)
        # 对图片数据进行预处理（ResNet特定的预处理）
        img_data = preprocess_input(img_data)
        # 使用模型进行特征提取
        features = model.predict(img_data, verbose=0)
        # 将特征数组展平并转换为Python列表返回
        return features.flatten().tolist()

    except Exception as e:
        print(f"特征提取错误进程ID: {os.getpid()} 文件路径: {img_path}: {str(e)}")
        return None


def res_nets(model, img_paths):
    """并行特征提取"""
    print(f"正在进行特征提取...")
    num_workers = int(multiprocessing.cpu_count() / 2)  # 考虑到其他进程的干扰，将进程数设为 CPU 数量的 1/2
    with multiprocessing.Pool(
            processes=num_workers,
            initializer=init_model,
            initargs=(model,)
    ) as pool:
        features = pool.map(partial(_extract_single_feature, model_type=model), img_paths)

    res_features = [f for f in features if f is not None]

    """ 独立进程导入模型 """
    from django.apps import apps
    print('保存特征...')
    Feature = apps.get_model('ipo', 'Feature')
    time.sleep(random.randrange(0, 10))
    Feature.objects.update_or_create(model=model, defaults={'model': model, 'features': res_features})
